import rosbag
import numpy as np
import sensor_msgs.point_cloud2 as pc2
import open3d as o3d
import cv2
import tqdm
from cv_bridge import CvBridge
from tictoc import TicToc
import csv
import os
import json
import argparse
import multiprocessing
from utils import topic2path, get_ms
from my_io import write_points, copy_file, files_filter


def main(args):
    # 统计耗时
    cost = TicToc("ros解析")
    assert os.path.exists(args.rosbag_dir)
    files = files_filter(args.rosbag_dir, pos_fix="bag", split='.')
    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args.rosbag_dir, args.data_path))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args.rosbag_dir, args.data_path)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(bag, rosbag_dir, data_path):
    rosbag_path = os.path.join(rosbag_dir, bag)
    dir_name = os.path.join(data_path, bag.split('.')[0])
    print("开始解析<{}>......".format(rosbag_path))

    bag = rosbag.Bag(rosbag_path)
    #
    lidar_topics = ["/pointcloud_lidar1", "/pointcloud_lidar2", "/pointcloud_lidar3", "/pointcloud_lidar4"]
    radar_topics = ["/pointcloud_radar3"]
    old_radar_topics = {"/radar_det_list": "/pointcloud_radar3"}
    # camera_topics = ["/camera71/compressed", "/camera72/compressed", "/camera73/compressed", "/camera74/compressed",
    #                  "/camera75/compressed", "/camera76/compressed",
    #                  "/camera77/compressed", "/camera78/compressed", "/camera79/compressed", "/camera80/compressed",
    #                  "/camera81/compressed", "/camera82/compressed"]
    camera_topics = ["/camera75/compressed", "/camera77/compressed", "/camera80/compressed",
                     "/camera82/compressed"]
                     # "/camera81/compressed"]
    camera_ori_topics = []

    # lidar_topics = ["/pointcloud_lidar1", "/pointcloud_lidar2", "/pointcloud_lidar3"]
    # radar_topics = ["/pointcloud_radar3"]
    # old_radar_topics = {}
    # camera_topics = []
    # camera_ori_topics = ["/camera71", "/camera72", "/camera73", "/camera74",
    #                  "/camera75", "/camera76",
    #                   "/camera79", "/camera80",
    #                  "/camera82"]

    localization_topics = ["/localization_result"]
    print("保存在<{}>".format(dir_name))

    # 新建目录
    new_dir = []
    ffarther_dir = []
    farther_dir = []
    child_dir = []
    ffarther_dir.append(dir_name)
    new_dir = new_dir + ffarther_dir
    farther_dir.append(os.path.join(dir_name, "lidar"))
    farther_dir.append(os.path.join(dir_name, "radar"))  # 毫米波所以直接保存为bin:因为1)不直接标注 2)pcd不能保留r_speed和rcs
    farther_dir.append(os.path.join(dir_name, "camera"))
    farther_dir.append(os.path.join(dir_name, "localization"))
    farther_dir.append(os.path.join(dir_name, "calib"))
    new_dir = new_dir + farther_dir
    for topic in lidar_topics:
        child_dir.append(os.path.join(dir_name, "lidar", topic2path(topic)))
    for topic in radar_topics:
        child_dir.append(os.path.join(dir_name, "radar", topic2path(topic)))
    for topic in camera_topics:
        child_dir.append(os.path.join(dir_name, "camera", topic2path(topic, -2)))
    for topic in camera_ori_topics:
        child_dir.append(os.path.join(dir_name, "camera", topic2path(topic)))
    new_dir = new_dir + child_dir
    for dir in new_dir:
        if not os.path.exists(dir):
            print("新建目录<{}>".format(dir))
            os.makedirs(dir, exist_ok=True)

    if os.path.exists(os.path.join(rosbag_dir, 'calib.json')):
        copy_file(os.path.join(rosbag_dir, 'calib.json'), os.path.join(dir_name, 'calib'))
        print("拷贝 calib.json , from {} to {}".format(os.path.join(rosbag_dir, 'calib'),
                                                       os.path.join(dir_name, 'calib')))
    else:
        # 创建calib.yml
        eye44 = np.eye(4).tolist()  # 4*4的旋转平移矩阵
        eye33 = np.eye(3).tolist()  # 3*3的旋转平移矩阵
        calib_file = os.path.join(dir_name, 'calib', 'calib.json')
        calib_data = dict()
        for topic in lidar_topics + radar_topics:
            sensor_id = topic[-6:]
            calib_data[sensor_id] = dict()
            calib_data[sensor_id]['transforms'] = dict()
            calib_data[sensor_id]['transforms']['base_link'] = eye44

        calib_data['camera75'] = dict()
        calib_data['camera75']['ori'] = dict()
        calib_data['camera75']['ori']['imgh'] = 9999
        calib_data['camera75']['ori']['imgw'] = 9999
        calib_data['camera75']['ori']['K'] = eye33
        calib_data['camera75']['ori']['D'] = np.zeros(shape=(0, 5)).tolist()

        # 写入JSON文件
        with open(calib_file, 'w', encoding='utf-8') as file:
            json.dump(calib_data, file, ensure_ascii=False, indent=4)

        # fs = cv2.FileStorage(os.path.join(dir_name, 'calib', 'sensor.yml'), cv2.FileStorage_WRITE)
        # matrix = np.eye(4)  # 4*4的旋转平移矩阵
        # for topic in lidar_topics + radar_topics:
        #     sensor_id = topic[-6:]
        #     fs.write(sensor_id, matrix)
        # fs.release()
        print("新建标定文件 in {}".format(calib_file))

    # 创建或打开文件
    csvfile = open(os.path.join(dir_name, "localization", 'localization.csv'), mode='w', newline='')
    # 标题列表
    fieldnames = ['stamp(ms)', 'latitude', 'longitude', 'altitude', 'mercator_x', 'mercator_y', 'mercator_z',
                  'vel_x', 'vel_y', 'vel_z', 'accel_x', 'accel_y', 'accel_z',
                  'roll', 'pitch', 'yaw', 'angular_x', 'angular_y', 'angular_z']
    # 创建 DictWriter 对象
    write = csv.DictWriter(csvfile, fieldnames=fieldnames)
    # 写入表头
    write.writeheader()

    for topic, msg, t in bag:
        if topic in lidar_topics:
            lidar = pc2.read_points(msg)
            points = np.array(list(lidar))
            # # 创建一个PointCloud对象
            # pcd = o3d.geometry.PointCloud()
            # # 将随机数转换成PointCloud点数据
            # pcd.points = o3d.utility.Vector3dVector(points[:, :3])
            # # 将PointCloud点数据保存成pcd文件，格式为assii文本格式 比如用于lc标定任务
            # o3d.io.write_point_cloud(os.path.join(dir_name,'lidar', topic2path(topic), str(get_ms(msg)) + ".pcd"), pcd,
            #                          write_ascii=True)
            assert points.shape[1] >= 4
            write_points(points[:, [0, 1, 2, 3, 5]],
                         os.path.join(dir_name, 'lidar', topic2path(topic), str(get_ms(msg)) + ".bin"))

        elif topic in radar_topics:
            radar = pc2.read_points(msg)
            points = np.array(list(radar))
            points = points[:, :5]  # x y z r_speed rcs
            file_path = os.path.join(dir_name, 'radar', topic2path(topic), str(get_ms(msg)) + ".bin")
            write_points(points, file_path)
            # # 创建一个PointCloud对象
            # pcd = o3d.geometry.PointCloud()
            # # 将随机数转换成PointCloud点数据
            # pcd.points = o3d.utility.Vector3dVector(points[:, :3])# x y z r_speed rcs
            # # 将PointCloud点数据保存成pcd文件，格式为assii文本格式
            # o3d.io.write_point_cloud(os.path.join(dir_name,'radar', topic2path(topic), str(get_ms(msg)) + ".pcd"), pcd,
            #                          write_ascii=True)

        elif topic in old_radar_topics:
            points_list = []
            for det in msg.detection_array:
                points_list.append([det.f_x, det.f_y, det.f_z, det.f_RangeRate, det.s_RCS])
            points = np.array(points_list)
            file_path = os.path.join(dir_name, 'radar', topic2path(old_radar_topics[topic]),
                                     str(get_ms(msg.detection_array[0])) + ".bin")
            write_points(points, file_path)

        elif topic in camera_topics:
            bridge = CvBridge()
            cv_image = bridge.compressed_imgmsg_to_cv2(msg, "bgr8")
            # tm4
            # if topic[:9] == '/camera81' or topic[:9] == '/camera80':  # 前后相机超前600ms，左右相机超前500ms
            #     cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 450) + ".jpg"),
            #                 cv_image)
            # else:
            #     cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 600) + ".jpg"),
            #                 cv_image)
            # kb qj101
            if topic[:9] == '/camera81' or topic[:9] == '/camera80':  # 前后相机超前600ms，左右相机超前500ms
                cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 650) + ".jpg"),
                            cv_image)
            else:
                cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 750) + ".jpg"),
                            cv_image)

            # cv2.imshow("Image window", cv_image)
            # cv2.waitKey(3)
        elif topic in camera_ori_topics:
            bridge = CvBridge()
            cv_image = bridge.imgmsg_to_cv2(msg, "bgr8")
            # cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg)+600) + ".jpg"), cv_image)
            if topic[:9] == '/camera81':  # 前后相机超前600ms，左右相机超前500ms
                cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 450) + ".jpg"),
                            cv_image)
            else:
                cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 600) + ".jpg"),
                            cv_image)

        elif topic in localization_topics:
            # +'\t'用于取消科学计数法
            write.writerow({'stamp(ms)': str(get_ms(msg)) + '\t', 'latitude': msg.latitude, 'longitude': msg.longitude,
                            'altitude': msg.altitude, 'mercator_x': msg.mercator_x, 'mercator_y': msg.mercator_y,
                            'mercator_z': msg.mercator_z, 'vel_x': msg.vel_x, 'vel_y': msg.vel_y, 'vel_z': msg.vel_z,
                            'accel_x': msg.accel_x,
                            'accel_y': msg.accel_y, 'accel_z': msg.accel_z, 'roll': msg.roll, 'pitch': msg.pitch,
                            'yaw': msg.yaw,
                            'angular_x': msg.angular_x, 'angular_y': msg.angular_y,
                            'angular_z': msg.angular_z})

    for dir in child_dir:
        print("文件夹:<{}> 包含文件<{}>个".format(dir.split('/')[-1], len(os.listdir(dir))))

    bag.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--rosbag-dir',
                        default='/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/calib/bags',
                        help='rosbag')
    parser.add_argument('--data-path',
                        default='/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/calib/datasets')
    args = parser.parse_args()

    main(args)
